[executorch] Propagate device metadata from partitioner result onto TensorSpecs by Gasoonjia · Pull Request #18078 · pytorch/executorch

Gasoonjia · 2026-03-10T23:21:49Z

Stack from ghstack (oldest at bottom):

[ET Device Support] Annotate device attributes of CUDA backend IO tensors cuda device #18080
[ET Device Support] Parse device info from serialized tensor in tensor_parser #18328
[ET Device Support] Propagate device info from TensorSpec into serialized Tensor #18079
-> [executorch] Propagate device metadata from partitioner result onto TensorSpecs #18078

Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph

The overall pipeline is:
a. Partitioner use compile_spec to determine which device the partitoned blob is runing on
b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device.

Differential Revision: D95842511

…ensorSpecs Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) [ghstack-poisoned]

pytorch-bot · 2026-03-10T23:21:53Z

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18078

📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❗ 1 Active SEVs

There are 1 currently active SEVs. If your PR is affected, please view them below:

Rolling out OSDC (ARC) runners on pull workflow for PyTorch trunk commits

❌ 1 New Failure, 2 Unrelated Failures

As of commit 4ae5949 with merge base 0bf3c51 ():

NEW FAILURE - The following job has failed:

pull / unittest-editable / macos / macos-job (gh)
export/tests/test_target_recipes.py::TestTargetRecipes::test_mv2_model

BROKEN TRUNK - The following jobs failed but were present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

pull / unittest / windows / windows-job (gh) (trunk failure)
##[error]The operation was canceled.
pull / unittest-editable / windows / windows-job (gh) (trunk failure)
##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

github-actions · 2026-03-10T23:22:28Z

This PR needs a `release notes:` label

If your change should be included in the release notes (i.e. would users of this library care about this change?), please use a label starting with release notes:. This helps us keep track and include your important work in the next release notes.

To add a label, you can comment to pytorchbot, for example
@pytorchbot label "release notes: none"

For more information, see
https://github.com/pytorch/pytorch/wiki/PyTorch-AutoLabel-Bot#why-categorize-for-release-notes-and-how-does-it-work.

digantdesai · 2026-03-11T02:47:26Z

+                if lowered_module is None:
+                    continue
+
+                result = _get_target_device_from_compile_specs(lowered_module)


This effectively assumes that we know the device 'name' AoT. In theory, we can have a multi-device delegate then the runtime might interpret this name differently and that can cause some confusion i.e cuda:0 device on Metal.

I am not sure about using generic names like 'gpu' but also not sure about following PyTorch's eager/jit style naming convention where you won't switch devices underneath.

May I have your suggestions on the executorch device name?

Currently we set up the device name AOT and intentionally decouple dour device attribute with pytorch/pytorch device concept; we created a enum in the etensor schema for all devices we are supporting right now. In this way we can support as much as device as we want.

For the situaton you mentioned, if other backend like vulken need its own gpu device, they should add a new one to the enum. We should avoid using generic names like 'gpu'.

Multi-device graph serialization will necessitate multiple graphs. We can maybe make an exception for input tensors, but for any intermediate the runtime needs to know what the device its loading intermediates onto.

Device is fixed at export aot. If you want to have some generic shader style lib where the gpu type is decided lazily then you will have to use a generic key like gpu.

…sult onto TensorSpecs" Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) [ghstack-poisoned]

… onto TensorSpecs Pull Request resolved: #18078 Annotate the delegate's input and output tensors as specific device type The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. ghstack-source-id: 352045003 @exported-using-ghexport Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/)

…sult onto TensorSpecs" Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) [ghstack-poisoned]

… onto TensorSpecs Pull Request resolved: #18078 Annotate the delegate's input and output tensors as specific device type The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device with correct device index. ghstack-source-id: 363318415 @exported-using-ghexport Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/)

…sult onto TensorSpecs" Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) [ghstack-poisoned]

JacobSzwejbka · 2026-04-09T16:35:23Z

+    return device_type, device_index
+
+
+def _get_lowered_module(


nit: this type of util really should be placed in a single spot. There are other things like this in the passes. Lets take it as a follow up to have claude just search for generic utils like this and centralize them

JacobSzwejbka · 2026-04-09T16:36:19Z

+    device_index: int = 0,
+) -> None:
+    """Set the device attribute on a TensorSpec."""
+    spec.device = device_type


Are these fields already in the TensorSpec class definition? Are they initialized to just cpu and 0?

Yes and yes.

JacobSzwejbka · 2026-04-09T16:45:30Z

+        for node in graph_module.graph.nodes:
+            if node.op == "call_function" and node.target == executorch_call_delegate:
+                lowered_module = _get_lowered_module(graph_module, node)
+                if lowered_module is None:


We should throw here no?

Let me throw here. I don't think it will be None.

JacobSzwejbka · 2026-04-09T16:45:51Z

+                    continue
+
+                result = _get_target_device_from_compile_specs(lowered_module)
+                if result is None:


Why does it not return cpu by default

The default value for every tensor is cpu:0. If backend author didn't set the device-related compile spec, every tensor should be remain default value and no need to reset as cpu

JacobSzwejbka · 2026-04-09T16:53:11Z

+
+        # Second pass: propagate device through getitem nodes that extract
+        # individual outputs from a delegate call.
+        for node in graph_module.graph.nodes:


Can we just do 1 pass. You can look at users of the delegate node to find the getitem nodes.

Yes we can but i feel like two passes are more structural: one is specific for delegate input and the other is for delegate output

…sult onto TensorSpecs" Add end-to-end device type annotation support from export to runtime. Currently we only support one device per graph The overall pipeline is: a. Partitioner use `compile_spec` to determine which device the partitoned blob is runing on b. after lowered partitioned graph to backend, the new-introed propagate_device_pass will annotate the input and output tensors of delegate blob as target device. Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) [ghstack-poisoned]

…ized Tensor (#18079) Stack from [ghstack](https://github.com/ezyang/ghstack) (oldest at bottom): * #18080 * #18328 * __->__ #18079 * #18078 Propagate device information from `TensorSpec.device` (set by `PropagateDevicePass`) to the serialized `schema.Tensor` in the emitted PTE file, to make runtime further aware of it. Differential Revision: [D95899706](https://our.internmc.facebook.com/intern/diff/D95899706/)

@Gasoonjia

…ensorSpecs (#18893) This PR was created by the merge bot to help merge the original PR into the main branch. ghstack PR number: #18078 by @Gasoonjia ^ Please use this as the source of truth for the PR details, comments, and reviews ghstack PR base: https://github.com/pytorch/executorch/tree/gh/gasoonjia/135/base ghstack PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/135/head Merge bot PR base: https://github.com/pytorch/executorch/tree/main Merge bot PR head: https://github.com/pytorch/executorch/tree/gh/gasoonjia/135/orig Differential Revision: [D95842511](https://our.internmc.facebook.com/intern/diff/D95842511/) @diff-train-skip-merge Co-authored-by: gasoonjia <gasoonjia@icloud.com>

Gasoonjia requested review from JacobSzwejbka and larryliu0820 as code owners March 10, 2026 23:21

meta-cla bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Mar 10, 2026

meta-codesync bot added fb-exported meta-exported labels Mar 10, 2026

digantdesai reviewed Mar 11, 2026

View reviewed changes

Gasoonjia mentioned this pull request Feb 18, 2026

[ET Device Support] DeviceAllocator interface and DeviceAllocatorRegistry #17535

Open

Gasoonjia added 2 commits March 12, 2026 19:49

Gasoonjia added 4 commits March 14, 2026 17:14

Gasoonjia added 2 commits March 19, 2026 11:43

This was referenced Mar 20, 2026

[ET Device Support] Device-aware memory planning: separate buffers per device type #18375

Open

add device attribute in ETensor #17161

Open

lucylq approved these changes Mar 24, 2026

View reviewed changes

This was referenced Apr 6, 2026

[ET Device Support] Define AOT device copy ops registry #18728

Open

[ET Device Support] Define et_copy runtime h2d and d2h copy ops #18729

Open

[ET Device Support] PropagateDevicePass inserts H2D/D2H copy ops at delegate boundaries #18730

Open

JacobSzwejbka reviewed Apr 9, 2026

View reviewed changes

Gasoonjia added 2 commits April 13, 2026 15:41

Gasoonjia merged commit cebbb1b into gh/gasoonjia/135/base Apr 14, 2026
160 of 164 checks passed

Gasoonjia deleted the gh/gasoonjia/135/head branch April 14, 2026 20:57

Gasoonjia temporarily deployed to cherry-pick-bot April 14, 2026 20:57 — with GitHub Actions Inactive

pytorchbot mentioned this pull request Apr 14, 2026

[executorch] Propagate device metadata from partitioner result onto TensorSpecs #18893

Merged

Conversation

Gasoonjia commented Mar 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

pytorch-bot bot commented Mar 10, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/executorch/18078

❗ 1 Active SEVs

❌ 1 New Failure, 2 Unrelated Failures

Uh oh!

github-actions bot commented Mar 10, 2026

This PR needs a release notes: label

Uh oh!

digantdesai Mar 11, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

Gasoonjia commented Mar 10, 2026 •

edited

Loading

pytorch-bot bot commented Mar 10, 2026 •

edited

Loading

This PR needs a `release notes:` label

digantdesai Mar 11, 2026 •

edited

Loading